import os

import numpy as np
import torch
from torch.utils.data import Dataset
from preprocess.transforms import RandomCrop, RandomFlip_LR, RandomFlip_UD, Center_Crop, Compose, Resize


class TrainDataset(Dataset):
    def __init__(self, crop_size=(32, 64, 64),image_dir='dataset/image/',mask_dir="dataset/mask/"):
        self.image_dir = image_dir
        self.mask_dir = mask_dir
        self.length = len(os.listdir(image_dir))-1 # dir includes .keep,so - 1 
        self.transforms = Compose([
            RandomCrop(crop_size),
            RandomFlip_LR(prob=0.5),
            RandomFlip_UD(prob=0.5),
            # RandomRotate()
        ])

    def __getitem__(self, index):
        ct_array = np.load(self.image_dir + str(index) + '.npy')
        seg_array = np.load(self.mask_dir + str(index) + '.npy')

        ct_array = ct_array / 255.0
        ct_array = ct_array.astype(np.float32)

        ct_array = torch.FloatTensor(ct_array).unsqueeze(0)
        seg_array = torch.FloatTensor(seg_array).unsqueeze(0)

        if self.transforms:
            ct_array, seg_array = self.transforms(ct_array, seg_array)

        return ct_array, seg_array.squeeze(0)

    def __len__(self):
        return self.length
